# -*- coding: utf-8 -*-
import os
import hashlib
import shutil
import time
import sys


# 设置系统字符集，防止写入log时出现错误
try: #python2
    reload(sys)
    sys.setdefaultencoding('utf-8')
except: #python3
    pass

# ======================================================================================
class Foo(object):
    def __init__(self, base_dir, exclude, log_path):
       self.base_dir = os.path.realpath(base_dir) # 即程序运行时用来存放备份文件等各类文件的文件夹
       self.special_dir = os.path.realpath(os.path.join(log_path, '.drops_' + str(base_dir).replace('/','_').replace('\\','_').replace(':','_')))
       self.bak_dir = os.path.join(self.special_dir, 'bak') # 备份文件存放路径
       self.log_dir = os.path.join(self.special_dir, 'log') # 日志存放路径
       self.webshell_dir = os.path.join(self.special_dir, 'webshell') # webshell存放路径，即新生成的文件存放路径
       self.difffile_dir = os.path.join(self.special_dir, 'difffile') # 被修改的文件存放路径
       self.oprfile_config_path = os.path.join(self.base_dir, '.oprfile.txt') # 配置文件，记录正操作的文件，状态：-1代表排除文件或目录，0代表正在操作，1代表操作完毕。用来保证开启文件监控时也可以修改文件。且0状态是有必要的，因为修改过程中放弃修改还能还原回来，但是-1状态变为1状态后不能得到最原始的文件
       self.exclude = list(exclude) # 排除的关键字，如果路径中包含此关键字则不会监控
       self.exclude.append(self.oprfile_config_path)
       self.exclude.append(self.special_dir)
       self.create_all_dir()
       self.origin_file_list = self.getfilelist(self.base_dir) # 原始文件路径列表
       self.file_md5_dict = self.getfilemd5dict(self.origin_file_list) # 原始文件的md5字典
       self.backup_file(self.origin_file_list)

    # 获取当前所有文件路径
    def getfilelist(self, base_dir):
        filelist = []
        for root, subdirs, files in os.walk(base_dir):
            for filepath in files:
                originalfile = os.path.realpath(os.path.join(root, filepath))
                if self.is_excluded_file(originalfile):
                    continue
                if 'win' in sys.platform.lower() or os.access(filepath, os.R_OK) and os.access(filepath, os.W_OK): #因为os.access只支持linux
                    filelist.append(originalfile)
        return filelist

    # 判断文件是不是被排除的文件
    def is_excluded_file(self, filepath):
        for ex in self.exclude:
            if ex in filepath:
                return True
        return False

    def create_all_dir(self):
        self.mkdir_p(self.special_dir)
        self.mkdir_p(self.bak_dir)
        self.mkdir_p(self.log_dir)
        self.mkdir_p(self.webshell_dir)
        self.mkdir_p(self.difffile_dir)

    # 计算文件MD5值
    def calcMD5(self, filepath):
        try:
            with open(filepath,'rb') as f:
                md5obj = hashlib.md5()
                md5obj.update(f.read())
                hash = md5obj.hexdigest()
                return hash
        # 文件MD5消失即为文件被删除，恢复文件
        except Exception as e:
            print('[!] MD5 calculate error:' + str(e))

    # 获取所有文件MD5
    def getfilemd5dict(self, filelist = []):
        filemd5dict = {}
        for ori_file in filelist:
            md5 = self.calcMD5(os.path.realpath(ori_file))
            if md5:
                filemd5dict[ori_file] = md5
        return filemd5dict

    # 备份所有文件
    def backup_file(self, filelist=[]):
        for filepath in filelist:
            relative = filepath[len(self.base_dir)+1:]
            newpath = os.path.join(self.bak_dir, relative)
            self.copyfile(filepath, newpath)

    # 得到备份文件应该放的位置
    def get_bakfile_path(self, filepath):
        relative = filepath[len(self.base_dir)+1:]
        newpath = os.path.join(self.bak_dir, relative)
        return newpath

    # 拷贝文件
    def copyfile(self, oldpath, newpath):
        file_dir = os.path.dirname(newpath)
        if os.path.isdir(file_dir) == False:
            self.mkdir_p(file_dir)
        try:
            shutil.copy2(oldpath, newpath)
        except Exception as e:
            print('[!] Copy file error:' + str(e))
    
    # 还原文件，即将拷贝的文件复制一份放到原位置，比如文件被删除或者文件被修改时会用到
    def recover_file(self, filepath):
        self.copyfile(self.get_bakfile_path(filepath), filepath)

    # 检查正在操作的文件，状态：-1是排除文件或目录，0是正在操作，1是操作完成。这个功能是用来保证开启文件监控时可以进行文件的操作
    # 当发现状态是0时，暂时不对文件进行处理，从源文件列表中删除它
    # 当发现状态是1时，检查文件如果存在，认为是修改或添加了文件；如果文件不存在认为是删除了文件
    def handle_oprfile(self):
        if not os.path.exists(self.oprfile_config_path):
            f = open(self.oprfile_config_path, 'w')
            f.close()
            return []
        f=open(self.oprfile_config_path,'r')
        items = f.readlines()
        f.close()
        oprfiles = []
        oprs = []
        for item in items:
            try: #python2
                opr = item.encode('gbk').split(',') #python2编解码比较乱,这里是为了适应中文文件名的情况，参考 https://blog.csdn.net/liuweiyuxiang/article/details/83831184
            except: #python3
                opr = item.split(',')
            if len(opr) > 1:
                status = opr[1].strip()
                if status == '':
                    status = '0'
            else:
                status = '0'
            filepath = opr[0].strip()
            if filepath == '':
                continue
            filepath = os.path.realpath(os.path.join(self.base_dir, filepath)) #为了格式化路径，脚本里所有文件路径都必须是\\xx\\xx\\xx的形式
            if status == '-1': # 排除出去 ------------------------------------------------
                self.exclude.append(filepath)
                i = 0
                while i < len(self.origin_file_list): #从原始文件列表中也要将指定文件或目录排除出去
                    if os.path.isfile(filepath):
                        if filepath in self.origin_file_list[i]:
                            self.origin_file_list.remove(self.origin_file_list[i])
                        else:
                            i += 1
                    elif os.path.isdir(filepath):
                        file_dir = filepath + os.path.sep
                        if file_dir in self.origin_file_list[i]:
                            self.origin_file_list.remove(self.origin_file_list[i])
                        else:
                            i += 1
                print(r'[*] \(^v^)/ Exclude file or directory success:' + filepath)
            elif status == '0': # 正在被操作 ---------------------------------------------
                # 考虑vim修改文件会产生swp文件，也添加进去
                swpname = os.path.basename(filepath) + '.swp'
                if swpname[0] != '.':
                    swpname = '.' + swpname
                oprfiles.append(os.path.join(os.path.dirname(filepath), swpname))
                oprfiles.append(filepath)
                oprs.append(item)
            elif status == '1': # 操作完成 -----------------------------------------------
                if os.path.isfile(filepath):
                    # 文件存在说明可能被修改或新建文件
                    self.backup_file([filepath])
                    self.origin_file_list.append(filepath)
                    self.file_md5_dict[filepath] = self.calcMD5(filepath)
                    print(r'[+] \(^v^)/ Modify or add file success:' + filepath)
                    if filepath in self.exclude: # 如果文件在排除目录里，要从排除目录中删除
                        self.exclude.remove(filepath)
                elif os.path.isdir(filepath):
                    # 说明要包含进来的是文件夹，将里面的文件都包含进来，并备份文件
                    if filepath in self.exclude:
                        self.exclude.remove(filepath)
                    file_list = self.getfilelist(filepath)
                    self.backup_file(file_list)
                    for f in file_list:
                        self.file_md5_dict[f] = self.calcMD5(f)
                        if f not in self.origin_file_list:
                            self.origin_file_list.append(f)
                    print(r'[+] \(^v^)/ Include directory success:' + filepath)
                else:
                    # 文件不存在说明文件被删除，则从源文件列表中也要删除它
                    if filepath in self.origin_file_list:
                        self.origin_file_list.remove(filepath)
                    print(r'[-] \(^v^)/ Delete file success:' + filepath)
        # 重新写文件，这里只写正在操作的文件
        f = open(self.oprfile_config_path, 'w')
        f.writelines(oprs)
        f.close()
        return oprfiles

    # 核心函数，检查文件变化
    def check(self):
        oprfiles = self.handle_oprfile()
        file_list = self.getfilelist(self.base_dir)

        # 逐文件比较
        # 1、如果文件在源文件列表里，对比md5看是否被修改，被修改就还原它
        # 2、如果文件不在源文件列表里，说明是上传的新文件，要删掉它
        # 3、如果某源文件检查不到文件存在，说明被删除了，要还原它
        # 如下为实现过程：
        for filepath in file_list:
            if filepath in oprfiles: #如果文件正在被操作先不要处理
                continue
            if filepath in self.origin_file_list:
                # 【1、检查文件是否被修改】这里说明文件再源文件列表中，对比md5值
                file_md5 = self.calcMD5(filepath)
                if file_md5 != self.file_md5_dict[filepath]:
                    # md5值不相等，尝试还原文件
                    try:
                        print('[*] The file was modified:' + filepath)
                        shutil.move(filepath, os.path.join(self.difffile_dir, os.path.basename(filepath) + '.txt'))
                        self.recover_file(filepath)
                        print('[+] File restored!')
                    except Exception as e:
                        print('[!] Failed to move file, "%s" suspected webshell, please handle it in time'%filepath)
                    try:
                        f = open(os.path.join(self.log_dir, 'log.txt'), 'a')
                        f.write('difference_file: ' + filepath + ' 时间: ' + time.ctime() + '\n')
                        f.close()
                    except Exception as e:
                        print('[!] Record failed: modified file:' + filepath)
                        pass
            else:
                # 【2、检查是否有webshell】这里说明文件不在源文件中，可能是上传的webshell文件，要删除它
                try:
                    print('[*] Check and kill suspected webshell uploaded files: ' + filepath)
                    shutil.move(filepath, os.path.join(self.webshell_dir, os.path.basename(filepath) + '.txt'))
                    print('[-] The newly uploaded file has been deleted!')
                except Exception as e:
                    print('[!] Failed to move file, "%s" suspected webshell, please handle it in time'%filepath)
                try:
                    f = open(os.path.join(self.log_dir, 'log.txt'), 'a')
                    f.write('new_file: ' + filepath + ' 时间: ' + str(time.ctime()) + '\n')
                    f.close()
                except Exception as e:
                    print('[!] Record failed: upload file:' + str(e))
        
        # 【3、检查文件是否被删除】下面再检查源文件列表中的每个文件是否都存在
        for filepath in self.origin_file_list:
            if filepath in oprfiles: #如果文件正在被操作先不要处理
                continue
            if not os.path.isfile(filepath):
                # 文件不存在说明文件可能被删除了，这时要还原文件
                print('[*] File deleted:' + filepath)
                self.recover_file(filepath)
                print('[+] Deleted files have been recovered!')
                try:
                    f = open(os.path.join(self.log_dir, 'log.txt'), 'a')
                    f.write('deleted_file: ' + filepath + ' 时间: ' + time.ctime() + '\n')
                    f.close()
                except Exception as e:
                    print('[!] Record failed: deleted file:' + filepath)
                    pass

    # 目录创建
    def mkdir_p(self, path):
        import errno
        try:
            os.makedirs(path)
        except OSError as exc:
            if exc.errno == errno.EEXIST and os.path.isdir(path):
                pass
            else: raise

# ====================================================================================================================================
Special_string = 'drops_log'  # 免死金牌
exclude = []
directory = []
period = 1 #秒数
log_path = ''
help_str = u'''
使用示例:
    python file_monitor.py -d /var/www/html1 /var/www/html2 --exclude keyword1 keyword2 -t 2 --log-path /tmp

-d          要监视的目录，可以同时指定多个
--exclude   如果目录名称或者文件名称中包含关键字，就不会被监视
-t          多少秒检测一次，默认为1
--log-path  临时文件存放目录，默认为各自网站目录，如果指定了目录则以指定为准
-h          查看帮助

开启文件监控后，若想增删改文件：
1、在网站根目录下的.oprfile.txt文件中写入文件名称（注意必须是相对于根目录的名称)，例如 dir1/abc.php
2、当文件修改完毕后，在文件名称后面继续输入“,1”，例如 dir1/abc.php,1

开启文件监控后，若想排除目录或文件：
在网站根目录下的.oprfile.txt文件中写入文件名称（注意必须是相对于根目录的名称)，并标注状态为-1，例如 dir1,-1
'''
# 处理参数
args = sys.argv[1:]
prop = ""
for a in args:
    if a[0] == "-":
        if a == "--exclude":
            prop = "exclude"
        elif a == "-d":
            prop = "directory"
        elif a == "-t":
            prop = "period"
        elif a == "--log-path":
            prop = "log-path"
        elif a == "-h":
            print(help_str)
            exit()
        else:
            prop = ""
    else:
        if prop == "exclude":
            exclude.append(a)
        elif prop == "directory":
            directory.append(a)
        elif prop == "period":
            period = a
        elif prop == "log-path":
            log_path = a

if __name__ == '__main__':
    print('---------In the continuous monitoring file------------')
    foos = []
    if len(directory) == 0: 
        directory = [os.getcwd()]
    for item in directory:
        if log_path == '':
            lp = item
        else:
            lp = log_path
        foos.append(Foo(item, exclude, lp))
    print('[*] All files have been backed up!')

    while True:
        for foo in foos:
            foo.check()
        time.sleep(period)